import cv2
import time
import base64
import numpy as np
from flask_socketio import emit
import mediapipe as mp

class BufferList:
    def __init__(self, buffer_time, default_value=0):
        # 初始化一个固定大小的缓冲区，用于存储历史数据
        self.buffer = [default_value] * buffer_time

    def push(self, value):
        # 将新值添加到缓冲区，移除最早的值
        self.buffer.pop(0)
        self.buffer.append(value)

    def max(self):
        # 返回缓冲区中的最大值
        return max(self.buffer)

    def min(self):
        # 返回缓冲区中的最小值，过滤掉 None 值
        return min(filter(lambda x: x is not None, self.buffer), default=0)

    def smooth_update(self, old_value, new_value, alpha=0.5):
        # 对新值和旧值进行平滑更新，控制变化幅度
        return alpha * new_value + (1 - alpha) * old_value

class JumpRopeDetector:
    def __init__(self):
        # 初始化计数器和状态变量
        self.count = 0
        self.dir = 0
        self.stop = False
        self.point_sd = 0
        self.point_history = []
        self.cy_max, self.cy_min = 100, 100
        self.flip_flag = 250
        self.mod = ""
        self.beginCount = False
        self.startTime = -1
        self.newCount = 0

         # 初始化多个缓冲区用于平滑处理
        self.buffers = {
            "center_y": BufferList(50),
            "center_y_up": BufferList(50),
            "center_y_down": BufferList(50),
            "center_y_flip": BufferList(50),
            "center_y_pref_flip": BufferList(50),
        }

        # 设置阈值参数
        self.thresholds = {
            "buffer_time": 50,
            "dy_ratio": 0.3,
            "up_ratio": 0.55,
            "down_ratio": 0.35,
            "flag_low": 150,
            "flag_high": 250,
        }

        # Mediapipe pose 模块初始化
        self.mp_pose = mp.solutions.pose
        self.mp_drawing = mp.solutions.drawing_utils
        self.hip_landmarks = [23, 24] # 髋关节的关键点索引
        self.shoulder_landmarks = [11, 12] # 肩关节的关键点索引

    def smooth_point(self, point):
        # 对关键点坐标进行平滑处理
        self.point_history.append(point)
        if len(self.point_history) > 5:
            self.point_history.pop(0)
        avg_point = {
            "x": sum(p["x"] for p in self.point_history) / len(self.point_history),
            "y": sum(p["y"] for p in self.point_history) / len(self.point_history)
        }
        return avg_point

    def extract_landmarks(self, results, landmarks_indices, image_width, image_height):
        # 提取指定关键点的坐标信息
        return [
            (lm.x * image_width, lm.y * image_height)
            for i, lm in enumerate(results.pose_landmarks.landmark)
            if i in landmarks_indices
        ]

    def calculate_center_y(self, hip_points, shoulder_points):
        # 计算髋部和肩部关键点的中心 y 坐标以及肩髋垂直距离
        cy_hip = int(np.mean([point[1] for point in hip_points]))
        cy_shoulder = int(np.mean([point[1] for point in shoulder_points]))
        return cy_hip, cy_hip - cy_shoulder

    def update_counters(self, cy, cy_shoulder_hip, cy_max, cy_min, flip_flag, thresholds):
        # 更新跳绳状态标志和计数器
        dy = cy_max - cy_min
        if dy > thresholds["dy_ratio"] * cy_shoulder_hip:
            if (
                cy > cy_max - thresholds["up_ratio"] * dy
                and flip_flag == thresholds["flag_low"]
            ):
                flip_flag = thresholds["flag_high"]
            elif (
                cy < cy_min + thresholds["down_ratio"] * dy
                and flip_flag == thresholds["flag_high"]
            ):
                flip_flag = thresholds["flag_low"]
        return flip_flag

    def draw_visualizations(self, image, cx, cy, count, image_width, image_height):
        # 在图像上绘制计数和中心点
        cv2.circle(image, (cx, cy), 5, (0, 0, 255), -1)
        cv2.putText(
            image,
            "centroid",
            (cx - 25, cy - 25),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 0, 255),
            1,
        )
        cv2.putText(
            image,
            f"count = {count}",
            (int(image_width * 0.5), int(image_height * 0.4)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 255, 0),
            1,
        )

    def detectVideo(self, videoPath):
        self.startTime = time.time()
        self.mod = "video"
        # 检测视频文件中的跳绳动作
        success = True
        if not videoPath:
            print("No video file found")
            success = False
            return
        
        cap = cv2.VideoCapture(videoPath)
        with self.mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
            while success and not self.stop:
                success, img = cap.read()
                if not success:
                    break

                h, w, c = img.shape
                scale = 640 / max(w, h)
                img = cv2.resize(img, (0, 0), None, scale, scale)

                img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                results = pose.process(img_rgb)

                if results.pose_landmarks:
                    hip_points = self.extract_landmarks(results, self.hip_landmarks, w, h)
                    shoulder_points = self.extract_landmarks(results, self.shoulder_landmarks, w, h)
                    cx = int(np.mean([point[0] for point in hip_points]))
                    cy, cy_shoulder_hip = self.calculate_center_y(hip_points, shoulder_points)
                else:
                    cx, cy, cy_shoulder_hip = 0, 0, 0

                self.buffers["center_y"].push(cy)
                self.cy_max = self.buffers["center_y"].smooth_update(self.cy_max, self.buffers["center_y"].max())
                self.buffers["center_y_up"].push(self.cy_max)
                self.cy_min = self.buffers["center_y"].smooth_update(self.cy_min, self.buffers["center_y"].min())
                self.buffers["center_y_down"].push(self.cy_min)

                prev_flip_flag = self.flip_flag
                self.flip_flag = self.update_counters(
                    cy, cy_shoulder_hip, self.cy_max, self.cy_min, self.flip_flag, self.thresholds
                )
                self.buffers["center_y_flip"].push(self.flip_flag)
                self.buffers["center_y_pref_flip"].push(prev_flip_flag)

                pinlv = -1
                if prev_flip_flag < self.flip_flag:
                    self.count += 1
                    now = time.time()
                    if now-self.startTime > 1:
                        pinlv = (self.count - self.newCount) / (now - self.startTime)
                        self.newCount = self.count
                        self.startTime = now
                        pinlv = round(pinlv, 2)


                self.draw_visualizations(img, cx, cy, self.count, w, h)
                
                # 绘制所有关节点和连线
                self.mp_drawing.draw_landmarks(
                    img, results.pose_landmarks, self.mp_pose.POSE_CONNECTIONS)

                _, buffer = cv2.imencode('.jpg', img)
                jpg_as_text = base64.b64encode(buffer).decode('utf-8')
                emit('image_data', {'image': jpg_as_text, 'count': str(int(self.count)), 'cy': str(pinlv)}, broadcast=True)

                # cv2.imshow("Image", img)
                # cv2.waitKey(1)
        emit('finish')

    def detectCamera(self):
        self.mod = "camera"
        # 检测摄像头中的跳绳动作
        cap = cv2.VideoCapture(0)
        with self.mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
            while True and not self.stop:
                success, img = cap.read()
                if not success:
                    break

                h, w, c = img.shape
                scale = 640 / max(w, h)
                img = cv2.resize(img, (0, 0), None, scale, scale)

                img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                results = pose.process(img_rgb)

                if results.pose_landmarks:
                    hip_points = self.extract_landmarks(results, self.hip_landmarks, w, h)
                    shoulder_points = self.extract_landmarks(results, self.shoulder_landmarks, w, h)
                    cx = int(np.mean([point[0] for point in hip_points]))
                    cy, cy_shoulder_hip = self.calculate_center_y(hip_points, shoulder_points)
                else:
                    cx, cy, cy_shoulder_hip = 0, 0, 0

                self.buffers["center_y"].push(cy)
                self.cy_max = self.buffers["center_y"].smooth_update(self.cy_max, self.buffers["center_y"].max())
                self.buffers["center_y_up"].push(self.cy_max)
                self.cy_min = self.buffers["center_y"].smooth_update(self.cy_min, self.buffers["center_y"].min())
                self.buffers["center_y_down"].push(self.cy_min)

                prev_flip_flag = self.flip_flag
                self.flip_flag = self.update_counters(
                    cy, cy_shoulder_hip, self.cy_max, self.cy_min, self.flip_flag, self.thresholds
                )
                self.buffers["center_y_flip"].push(self.flip_flag)
                self.buffers["center_y_pref_flip"].push(prev_flip_flag)
                
                pinlv = -1

                if prev_flip_flag < self.flip_flag and self.beginCount:
                    if self.startTime == -1:
                        self.startTime = time.time()
                    
                    now = time.time()
                    self.count += 1
                    if now-self.startTime > 1:
                        pinlv = (self.count - self.newCount) / (now - self.startTime)
                        self.newCount = self.count
                        self.startTime = now
                        pinlv = round(pinlv, 2)

                    

                self.draw_visualizations(img, cx, cy, self.count, w, h)

                # 绘制所有关节点和连线
                self.mp_drawing.draw_landmarks(
                    img, results.pose_landmarks, self.mp_pose.POSE_CONNECTIONS)

                _, buffer = cv2.imencode('.jpg', img)
                jpg_as_text = base64.b64encode(buffer).decode('utf-8')
                # 计算每秒钟的跳绳次数
          
                emit('image_data', {'image': jpg_as_text, 'count': str(int(self.count)), 'cy': str(pinlv)}, broadcast=True)
         

                # cv2.imshow("Image", img)
                # cv2.waitKey(1)

    def stopTask(self):
        self.stop = True
        print("结束任务")
        emit('finish')

    def startCount(self):
        self.beginCount = True
        print("开始计数")

    def stopCount(self):
        self.beginCount = False
        self.startTime = -1
        self.count = 0
        print("停止计数")